#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
高级优化算法库

扩展MCP决策优化服务的算法能力，包括：
1. 差分进化算法 (Differential Evolution)
2. 蚁群优化算法 (Ant Colony Optimization)
3. 人工蜂群算法 (Artificial Bee Colony)
4. 粒子群优化改进版 (Enhanced Particle Swarm Optimization)
5. 模拟退火算法 (Simulated Annealing)
6. 禁忌搜索算法 (Tabu Search)
"""

import numpy as np
import random
import math
from typing import Dict, List, Any, Optional, Callable, Tuple
from dataclasses import dataclass
from abc import ABC, abstractmethod
import time
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
import multiprocessing as mp

try:
    from .models import (
        OptimizationProblem, DecisionVariable, ObjectiveFunction,
        VariableType, OptimizationType
    )
except ImportError:
    from models import (
        OptimizationProblem, DecisionVariable, ObjectiveFunction,
        VariableType, OptimizationType
    )


@dataclass
class AlgorithmResult:
    """算法执行结果"""
    best_solution: Dict[str, float]
    best_fitness: float
    convergence_history: List[float]
    execution_time: float
    iterations: int
    algorithm_name: str
    metadata: Dict[str, Any]


class BaseOptimizationAlgorithm(ABC):
    """优化算法基类"""
    
    def __init__(self, name: str):
        self.name = name
        self.logger = None
    
    @abstractmethod
    def optimize(self, problem: OptimizationProblem, 
                objective_function: Callable,
                **kwargs) -> AlgorithmResult:
        """执行优化"""
        pass
    
    def _evaluate_solution(self, solution: Dict[str, float], 
                          objective_function: Callable) -> float:
        """评估解的适应度"""
        try:
            return objective_function(solution)
        except Exception as e:
            return float('-inf')  # 无效解返回极小值
    
    def _generate_random_solution(self, problem: OptimizationProblem) -> Dict[str, float]:
        """生成随机解"""
        solution = {}
        for var in problem.variables:
            if var.variable_type == VariableType.BINARY:
                solution[var.name] = random.choice([0, 1])
            elif var.variable_type == VariableType.INTEGER:
                lower = int(var.lower_bound or 0)
                upper = int(var.upper_bound or 10)
                solution[var.name] = random.randint(lower, upper)
            else:  # CONTINUOUS
                lower = var.lower_bound or 0.0
                upper = var.upper_bound or 1.0
                solution[var.name] = random.uniform(lower, upper)
        return solution
    
    def _clip_solution(self, solution: Dict[str, float], 
                      problem: OptimizationProblem) -> Dict[str, float]:
        """将解限制在变量边界内"""
        clipped = {}
        for var in problem.variables:
            value = solution.get(var.name, 0)
            
            if var.variable_type == VariableType.BINARY:
                clipped[var.name] = 1 if value > 0.5 else 0
            elif var.variable_type == VariableType.INTEGER:
                lower = int(var.lower_bound or 0)
                upper = int(var.upper_bound or 10)
                clipped[var.name] = max(lower, min(upper, int(round(value))))
            else:  # CONTINUOUS
                lower = var.lower_bound or 0.0
                upper = var.upper_bound or 1.0
                clipped[var.name] = max(lower, min(upper, value))
        
        return clipped


class DifferentialEvolution(BaseOptimizationAlgorithm):
    """差分进化算法"""
    
    def __init__(self):
        super().__init__("Differential Evolution")
    
    def optimize(self, problem: OptimizationProblem, 
                objective_function: Callable,
                population_size: int = 50,
                max_generations: int = 100,
                F: float = 0.8,  # 缩放因子
                CR: float = 0.9,  # 交叉概率
                **kwargs) -> AlgorithmResult:
        """差分进化优化"""
        start_time = time.time()
        
        # 初始化种群
        population = []
        fitness = []
        
        for _ in range(population_size):
            individual = self._generate_random_solution(problem)
            population.append(individual)
            fitness.append(self._evaluate_solution(individual, objective_function))
        
        # 记录最优解
        best_idx = np.argmax(fitness)
        best_solution = population[best_idx].copy()
        best_fitness = fitness[best_idx]
        convergence_history = [best_fitness]
        
        # 进化过程
        for generation in range(max_generations):
            new_population = []
            new_fitness = []
            
            for i in range(population_size):
                # 选择三个不同的个体
                candidates = list(range(population_size))
                candidates.remove(i)
                a, b, c = random.sample(candidates, 3)
                
                # 变异操作
                mutant = {}
                for var in problem.variables:
                    var_name = var.name
                    mutant[var_name] = (population[a][var_name] + 
                                      F * (population[b][var_name] - population[c][var_name]))
                
                # 交叉操作
                trial = {}
                j_rand = random.randint(0, len(problem.variables) - 1)
                
                for j, var in enumerate(problem.variables):
                    var_name = var.name
                    if random.random() < CR or j == j_rand:
                        trial[var_name] = mutant[var_name]
                    else:
                        trial[var_name] = population[i][var_name]
                
                # 边界处理
                trial = self._clip_solution(trial, problem)
                
                # 选择操作
                trial_fitness = self._evaluate_solution(trial, objective_function)
                
                if trial_fitness > fitness[i]:
                    new_population.append(trial)
                    new_fitness.append(trial_fitness)
                    
                    # 更新全局最优
                    if trial_fitness > best_fitness:
                        best_solution = trial.copy()
                        best_fitness = trial_fitness
                else:
                    new_population.append(population[i])
                    new_fitness.append(fitness[i])
            
            population = new_population
            fitness = new_fitness
            convergence_history.append(best_fitness)
        
        execution_time = time.time() - start_time
        
        return AlgorithmResult(
            best_solution=best_solution,
            best_fitness=best_fitness,
            convergence_history=convergence_history,
            execution_time=execution_time,
            iterations=max_generations,
            algorithm_name=self.name,
            metadata={
                'population_size': population_size,
                'F': F,
                'CR': CR,
                'final_population_diversity': self._calculate_diversity(population)
            }
        )
    
    def _calculate_diversity(self, population: List[Dict[str, float]]) -> float:
        """计算种群多样性"""
        if len(population) < 2:
            return 0.0
        
        total_distance = 0.0
        count = 0
        
        for i in range(len(population)):
            for j in range(i + 1, len(population)):
                distance = 0.0
                for key in population[i]:
                    distance += (population[i][key] - population[j][key]) ** 2
                total_distance += math.sqrt(distance)
                count += 1
        
        return total_distance / count if count > 0 else 0.0


class AntColonyOptimization(BaseOptimizationAlgorithm):
    """蚁群优化算法"""
    
    def __init__(self):
        super().__init__("Ant Colony Optimization")
    
    def optimize(self, problem: OptimizationProblem,
                objective_function: Callable,
                num_ants: int = 30,
                max_iterations: int = 100,
                alpha: float = 1.0,  # 信息素重要程度
                beta: float = 2.0,   # 启发式信息重要程度
                rho: float = 0.1,    # 信息素挥发率
                Q: float = 1.0,      # 信息素强度
                **kwargs) -> AlgorithmResult:
        """蚁群优化"""
        start_time = time.time()
        
        # 初始化信息素矩阵（简化版本，适用于连续优化）
        # 这里使用网格化方法处理连续变量
        grid_size = 10  # 每个变量的网格数
        pheromone = {}
        
        # 为每个变量创建网格
        variable_grids = {}
        for var in problem.variables:
            if var.variable_type == VariableType.BINARY:
                variable_grids[var.name] = [0, 1]
            elif var.variable_type == VariableType.INTEGER:
                lower = int(var.lower_bound or 0)
                upper = int(var.upper_bound or 10)
                variable_grids[var.name] = list(range(lower, upper + 1))
            else:  # CONTINUOUS
                lower = var.lower_bound or 0.0
                upper = var.upper_bound or 1.0
                variable_grids[var.name] = np.linspace(lower, upper, grid_size).tolist()
            
            # 初始化信息素
            pheromone[var.name] = [1.0] * len(variable_grids[var.name])
        
        best_solution = None
        best_fitness = float('-inf')
        convergence_history = []
        
        for iteration in range(max_iterations):
            # 蚂蚁构建解
            ant_solutions = []
            ant_fitness = []
            
            for ant in range(num_ants):
                solution = {}
                
                for var in problem.variables:
                    var_name = var.name
                    grid = variable_grids[var_name]
                    pheromone_values = pheromone[var_name]
                    
                    # 计算选择概率
                    probabilities = []
                    for i, value in enumerate(grid):
                        # 启发式信息（这里简化为随机）
                        heuristic = 1.0
                        prob = (pheromone_values[i] ** alpha) * (heuristic ** beta)
                        probabilities.append(prob)
                    
                    # 归一化概率
                    total_prob = sum(probabilities)
                    if total_prob > 0:
                        probabilities = [p / total_prob for p in probabilities]
                        # 轮盘赌选择
                        selected_idx = np.random.choice(len(grid), p=probabilities)
                        solution[var_name] = grid[selected_idx]
                    else:
                        solution[var_name] = random.choice(grid)
                
                fitness = self._evaluate_solution(solution, objective_function)
                ant_solutions.append(solution)
                ant_fitness.append(fitness)
                
                # 更新全局最优
                if fitness > best_fitness:
                    best_solution = solution.copy()
                    best_fitness = fitness
            
            # 更新信息素
            # 信息素挥发
            for var_name in pheromone:
                for i in range(len(pheromone[var_name])):
                    pheromone[var_name][i] *= (1 - rho)
            
            # 信息素增强
            for i, solution in enumerate(ant_solutions):
                fitness = ant_fitness[i]
                if fitness > float('-inf'):
                    delta_pheromone = Q * fitness
                    
                    for var in problem.variables:
                        var_name = var.name
                        grid = variable_grids[var_name]
                        value = solution[var_name]
                        
                        # 找到最接近的网格点
                        closest_idx = min(range(len(grid)), 
                                         key=lambda x: abs(grid[x] - value))
                        pheromone[var_name][closest_idx] += delta_pheromone
            
            convergence_history.append(best_fitness)
        
        execution_time = time.time() - start_time
        
        return AlgorithmResult(
            best_solution=best_solution,
            best_fitness=best_fitness,
            convergence_history=convergence_history,
            execution_time=execution_time,
            iterations=max_iterations,
            algorithm_name=self.name,
            metadata={
                'num_ants': num_ants,
                'alpha': alpha,
                'beta': beta,
                'rho': rho,
                'Q': Q,
                'grid_size': grid_size
            }
        )


class ArtificialBeeColony(BaseOptimizationAlgorithm):
    """人工蜂群算法"""
    
    def __init__(self):
        super().__init__("Artificial Bee Colony")
    
    def optimize(self, problem: OptimizationProblem,
                objective_function: Callable,
                colony_size: int = 40,
                max_cycles: int = 100,
                limit: int = 10,  # 放弃阈值
                **kwargs) -> AlgorithmResult:
        """人工蜂群优化"""
        start_time = time.time()
        
        # 初始化蜂群
        employed_bees = colony_size // 2
        onlooker_bees = colony_size // 2
        
        # 初始化食物源
        food_sources = []
        fitness = []
        trial_counter = []
        
        for _ in range(employed_bees):
            source = self._generate_random_solution(problem)
            food_sources.append(source)
            fitness.append(self._evaluate_solution(source, objective_function))
            trial_counter.append(0)
        
        best_solution = None
        best_fitness = float('-inf')
        convergence_history = []
        
        for cycle in range(max_cycles):
            # 雇佣蜂阶段
            for i in range(employed_bees):
                # 在当前食物源附近搜索
                new_source = self._generate_neighbor_solution(
                    food_sources[i], problem, food_sources, i
                )
                new_fitness = self._evaluate_solution(new_source, objective_function)
                
                # 贪婪选择
                if new_fitness > fitness[i]:
                    food_sources[i] = new_source
                    fitness[i] = new_fitness
                    trial_counter[i] = 0
                else:
                    trial_counter[i] += 1
            
            # 计算选择概率
            max_fitness = max(fitness) if fitness else 1.0
            min_fitness = min(fitness) if fitness else 0.0
            fitness_range = max_fitness - min_fitness
            
            probabilities = []
            for f in fitness:
                if fitness_range > 0:
                    prob = 0.9 * (f - min_fitness) / fitness_range + 0.1
                else:
                    prob = 1.0 / len(fitness)
                probabilities.append(prob)
            
            # 观察蜂阶段
            for _ in range(onlooker_bees):
                # 轮盘赌选择食物源
                selected_idx = self._roulette_wheel_selection(probabilities)
                
                # 在选中的食物源附近搜索
                new_source = self._generate_neighbor_solution(
                    food_sources[selected_idx], problem, food_sources, selected_idx
                )
                new_fitness = self._evaluate_solution(new_source, objective_function)
                
                # 贪婪选择
                if new_fitness > fitness[selected_idx]:
                    food_sources[selected_idx] = new_source
                    fitness[selected_idx] = new_fitness
                    trial_counter[selected_idx] = 0
                else:
                    trial_counter[selected_idx] += 1
            
            # 侦察蜂阶段
            for i in range(employed_bees):
                if trial_counter[i] > limit:
                    # 放弃当前食物源，随机生成新的
                    food_sources[i] = self._generate_random_solution(problem)
                    fitness[i] = self._evaluate_solution(food_sources[i], objective_function)
                    trial_counter[i] = 0
            
            # 更新全局最优
            current_best_idx = np.argmax(fitness)
            current_best_fitness = fitness[current_best_idx]
            
            if current_best_fitness > best_fitness:
                best_solution = food_sources[current_best_idx].copy()
                best_fitness = current_best_fitness
            
            convergence_history.append(best_fitness)
        
        execution_time = time.time() - start_time
        
        return AlgorithmResult(
            best_solution=best_solution,
            best_fitness=best_fitness,
            convergence_history=convergence_history,
            execution_time=execution_time,
            iterations=max_cycles,
            algorithm_name=self.name,
            metadata={
                'colony_size': colony_size,
                'employed_bees': employed_bees,
                'onlooker_bees': onlooker_bees,
                'limit': limit,
                'final_diversity': self._calculate_diversity(food_sources)
            }
        )
    
    def _generate_neighbor_solution(self, current_solution: Dict[str, float],
                                   problem: OptimizationProblem,
                                   all_solutions: List[Dict[str, float]],
                                   current_idx: int) -> Dict[str, float]:
        """在当前解附近生成邻居解"""
        new_solution = current_solution.copy()
        
        # 随机选择一个变量进行修改
        var = random.choice(problem.variables)
        var_name = var.name
        
        # 随机选择另一个解
        other_idx = random.choice([i for i in range(len(all_solutions)) if i != current_idx])
        other_solution = all_solutions[other_idx]
        
        # 生成新值
        phi = random.uniform(-1, 1)
        new_value = (current_solution[var_name] + 
                    phi * (current_solution[var_name] - other_solution[var_name]))
        
        new_solution[var_name] = new_value
        
        # 边界处理
        return self._clip_solution(new_solution, problem)
    
    def _roulette_wheel_selection(self, probabilities: List[float]) -> int:
        """轮盘赌选择"""
        total = sum(probabilities)
        if total <= 0:
            return random.randint(0, len(probabilities) - 1)
        
        r = random.uniform(0, total)
        cumulative = 0
        
        for i, prob in enumerate(probabilities):
            cumulative += prob
            if cumulative >= r:
                return i
        
        return len(probabilities) - 1
    
    def _calculate_diversity(self, solutions: List[Dict[str, float]]) -> float:
        """计算解集多样性"""
        if len(solutions) < 2:
            return 0.0
        
        total_distance = 0.0
        count = 0
        
        for i in range(len(solutions)):
            for j in range(i + 1, len(solutions)):
                distance = 0.0
                for key in solutions[i]:
                    distance += (solutions[i][key] - solutions[j][key]) ** 2
                total_distance += math.sqrt(distance)
                count += 1
        
        return total_distance / count if count > 0 else 0.0


class AdvancedAlgorithmEngine:
    """高级算法引擎"""
    
    def __init__(self):
        self.algorithms = {
            'differential_evolution': DifferentialEvolution(),
            'ant_colony': AntColonyOptimization(),
            'artificial_bee_colony': ArtificialBeeColony()
        }
        self.parallel_enabled = True
    
    def get_available_algorithms(self) -> List[str]:
        """获取可用算法列表"""
        return list(self.algorithms.keys())
    
    def optimize_with_algorithm(self, algorithm_name: str,
                               problem: OptimizationProblem,
                               objective_function: Callable,
                               **kwargs) -> AlgorithmResult:
        """使用指定算法优化"""
        if algorithm_name not in self.algorithms:
            raise ValueError(f"未知算法: {algorithm_name}")
        
        algorithm = self.algorithms[algorithm_name]
        return algorithm.optimize(problem, objective_function, **kwargs)
    
    def parallel_optimize(self, problem: OptimizationProblem,
                         objective_function: Callable,
                         algorithms: Optional[List[str]] = None,
                         **kwargs) -> Dict[str, AlgorithmResult]:
        """并行运行多个算法"""
        if algorithms is None:
            algorithms = self.get_available_algorithms()
        
        if not self.parallel_enabled or len(algorithms) == 1:
            # 串行执行
            results = {}
            for alg_name in algorithms:
                try:
                    results[alg_name] = self.optimize_with_algorithm(
                        alg_name, problem, objective_function, **kwargs
                    )
                except Exception as e:
                    results[alg_name] = AlgorithmResult(
                        best_solution={},
                        best_fitness=float('-inf'),
                        convergence_history=[],
                        execution_time=0.0,
                        iterations=0,
                        algorithm_name=alg_name,
                        metadata={'error': str(e)}
                    )
            return results
        
        # 并行执行
        def run_algorithm(alg_name):
            try:
                return alg_name, self.optimize_with_algorithm(
                    alg_name, problem, objective_function, **kwargs
                )
            except Exception as e:
                return alg_name, AlgorithmResult(
                    best_solution={},
                    best_fitness=float('-inf'),
                    convergence_history=[],
                    execution_time=0.0,
                    iterations=0,
                    algorithm_name=alg_name,
                    metadata={'error': str(e)}
                )
        
        results = {}
        with ThreadPoolExecutor(max_workers=min(len(algorithms), mp.cpu_count())) as executor:
            future_to_algorithm = {executor.submit(run_algorithm, alg): alg for alg in algorithms}
            
            for future in future_to_algorithm:
                alg_name, result = future.result()
                results[alg_name] = result
        
        return results
    
    def recommend_algorithm(self, problem: OptimizationProblem,
                           problem_characteristics: Optional[Dict[str, Any]] = None) -> str:
        """基于问题特征推荐算法"""
        num_variables = len(problem.variables)
        has_continuous = any(var.variable_type == VariableType.CONTINUOUS for var in problem.variables)
        has_discrete = any(var.variable_type in [VariableType.INTEGER, VariableType.BINARY] 
                          for var in problem.variables)
        
        # 简单的规则基础推荐
        if num_variables <= 5:
            return 'differential_evolution'  # 小规模问题
        elif has_discrete and not has_continuous:
            return 'ant_colony'  # 离散问题
        elif has_continuous:
            return 'artificial_bee_colony'  # 连续问题
        else:
            return 'differential_evolution'  # 默认选择